import os
import pickle
from time import strftime

def save(fname):
    '用于记录收入'
    # 取出当前时间
    date = strftime('%Y-%m-%d')
    # 获取金额和说明
    try:
        amount = int(input('金额: '))
        comment = input('备注: ')
    except (KeyboardInterrupt, EOFError):
        print('\nBye-bye')
        exit(1)   # exit结束程序
    except ValueError:
        print('无效的金额。')
        return   # 函数的return可以提前结束函数
    # 取出收支记录
    with open(fname, 'rb') as fobj:
        records = pickle.load(fobj)
    # 计算最新余额
    balance = records[-1][-2] + amount
    # 构建一行记录，写入到收支记录列表中
    line = [date, amount, 0, balance, comment]
    records.append(line)
    # 存盘
    with open(fname, 'wb') as fobj:
        pickle.dump(records, fobj)

def cost(fname):
    '用于记录支出'
    # 取出当前时间
    date = strftime('%Y-%m-%d')
    # 获取金额和说明
    try:
        amount = int(input('金额: '))
        comment = input('备注: ')
    except (KeyboardInterrupt, EOFError):
        print('\nBye-bye')
        exit(1)   # exit结束程序
    except ValueError:
        print('无效的金额。')
        return   # 函数的return可以提前结束函数
    # 取出收支记录
    with open(fname, 'rb') as fobj:
        records = pickle.load(fobj)
    # 计算最新余额
    balance = records[-1][-2] - amount
    # 构建一行记录，写入到收支记录列表中
    line = [date, 0, amount, balance, comment]
    records.append(line)
    # 存盘
    with open(fname, 'wb') as fobj:
        pickle.dump(records, fobj)

def query(fname):
    '用于查询收支情况'
    # 取出数据
    with open(fname, 'rb') as fobj:
        records = pickle.load(fobj)
    # 打印表头
    print(f'{"date":<15}{"save":<8}{"cost":<8}{"balance":<10}{"comment":<20}')
    # 打印记录
    for line in records:
        print(f'{line[0]:<15}{line[1]:<8}{line[2]:<8}{line[3]:<10}{line[4]:<20}')

def show_menu():
    '显示主菜单'
    funcs = {'0': save, '1': cost, '2': query}
    prompt = '''(0) 收入
(1) 支出
(2) 查询
(3) 退出
请选择(0/1/2/3): '''
    fname = "money.data"  # 定义用于记账的文件
    # 记账文件如果不存在，则初始化它
    if not os.path.exists(fname):
        records = [[strftime('%Y-%m-%d'), 0, 0, 10000, 'init data']]
        with open(fname, 'wb') as fobj:
            pickle.dump(records, fobj)

    while 1:
        try:
            choice = input(prompt).strip()
        except (KeyboardInterrupt, EOFError):
            choice = '3'  # 如果用户按ctrl+c或ctrl+d，算他选了3

        if choice not in ['0', '1', '2', '3']:
            print('无效的输入，请重试。')
            continue

        if choice == '3':
            print('\nBye-bye')
            break

        funcs[choice](fname)

if __name__ == '__main__':
    show_menu()
